"""
Author: Morphlng
Date: 2024-04-15 14:03:13
LastEditTime: 2024-04-16 16:09:01
LastEditors: Morphlng
Description: Stub adapter for testing api without simulation
FilePath: /DrivingGym/src/driving_gym/simulation/adapter/stub/stub_adapter.py
"""

from __future__ import annotations

import logging

import numpy as np

from driving_gym.misc.util import override
from driving_gym.simulation.adapter_interface import AdapterInterface
from driving_gym.simulation.common.geom import (
    BoundingBox,
    Location,
    Transform,
    Vector3D,
)
from driving_gym.simulation.common.hd_map import HDMap

logger = logging.getLogger(__name__)


class StubAdapter(AdapterInterface):
    def __init__(self, config: dict = None):
        super().__init__(config)
        self._hd_map = None

        self._actors = {}
        self._sensors = {}
        self._callbacks = []
        self._frame = 0

    @property
    def hd_map(self):
        if self._hd_map is None:
            self._hd_map = HDMap(
                '<?xml version="1.0" encoding="utf-8"?><OpenDRIVE xmlns:xsd="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"><header revMajor="1" revMinor="4" name="" version="1" date="Thu, Aug 27 13:21:25 2020" north="0" south="0" east="0" west="0" vendor="LGSVL"><geoReference>+proj=tmerc +lat_0=37.4169056333992 +lon_0=-122.016036660455 +k=1 +x_0=587070 +y_0=4141577 +datum=WGS84 +units=m +no_defs </geoReference></header><road name="" length="251.00004577636719" id="0" junction="-1"><link /><type s="0" type="town"><speed max="25.0000406378746" unit="mph" /></type><planView><geometry s="0" x="125.50001525878906" y="0.04572296142578125" hdg="-3.1412472724914551" length="125.50002288818359"><line /></geometry><geometry s="125.50002288818359" x="0" y="0" hdg="-3.1412472724914551" length="125.50002288818359"><line /></geometry></planView><elevationProfile><elevation s="0" a="0" b="0" c="0" d="0" /><elevation s="125.50002288818359" a="0" b="0" c="0" d="0" /></elevationProfile><lateralProfile /><lanes><laneSection s="0"><left><lane id="1" type="driving" level="false"><link /><width sOffset="0" a="3.6287155151367188" b="-0.00022883601195644587" c="0" d="0" /><roadMark sOffset="0" type="solid" weight="standard" color="standard" laneChange="none" /></lane></left><center><lane id="0" type="driving" level="false"><link /><roadMark sOffset="0" type="solid" weight="standard" color="standard" laneChange="none" /></lane></center><right><lane id="-1" type="driving" level="false"><link /><width sOffset="0" a="3.534271240234375" b="0.00052373763173818588" c="0" d="0" /><roadMark sOffset="0" type="solid" weight="standard" color="standard" laneChange="none" /></lane></right></laneSection></lanes><objects /><signals /></road></OpenDRIVE>'
            )
        return self._hd_map

    @override(AdapterInterface)
    def initialize(self):
        return None

    @override(AdapterInterface)
    def cleanup(self):
        return None

    @override(AdapterInterface)
    def get_map(self) -> HDMap:
        return self.hd_map

    @override(AdapterInterface)
    def get_transform(self, identifier: str):
        return Transform()

    @override(AdapterInterface)
    def get_velocity(self, identifier: str):
        return Vector3D()

    @override(AdapterInterface)
    def get_acceleration(self, identifier: str):
        return Vector3D()

    @override(AdapterInterface)
    def get_bounding_box(self, identifier: str):
        return BoundingBox(Location(), Vector3D(2.4, 1.2, 0.8))

    @override(AdapterInterface)
    def get_sensor_data(self, sensor_identifier: str):
        sensor_config = self._sensors[sensor_identifier]
        shape = sensor_config["shape"]
        return np.random.rand(*shape)

    @override(AdapterInterface)
    def get_actors(self):
        return list(self._actors.keys())

    @override(AdapterInterface)
    def get_sensors(self):
        return list(self._sensors.keys())

    @override(AdapterInterface)
    def load_scenario(self, scenario: dict):
        actors = scenario["actor"]
        self._actors.update(actors)

        for actor_name, actor_config in actors.items():
            sensors = actor_config.get("sensors", {})
            for sensor_name, sensor_config in sensors.items():
                self._sensors[f"{actor_name}:{sensor_name}"] = sensor_config

    @override(AdapterInterface)
    def register_callback(self, func):
        self._callbacks.append(func)

    @override(AdapterInterface)
    def apply_control(self, target: str, control: dict):
        logger.info(f"{target}'s control: {control}")

    @override(AdapterInterface)
    def tick(self) -> int:
        for func in self._callbacks:
            func()

        self._frame += 1
        return self._frame


class StubSensorConfig:
    def __init__(self, sensor_name: str = "pseudo_sensor"):
        self.name = sensor_name
        self.shape = (168, 168, 3)

    def shaping(self, shape: tuple):
        self.shape = shape
        return self

    def to_dict(self):
        return {"name": self.name, "shape": self.shape}


class StubActorConfig:
    def __init__(self, actor_name: str):
        self.name = actor_name
        self.sensors = {}

    def add_sensor(self, sensor_name: str, sensor_config: "dict | StubSensorConfig"):
        if not isinstance(sensor_config, dict):
            sensor_config = sensor_config.to_dict()

        self.sensors[sensor_name] = sensor_config
        return self

    def to_dict(self):
        return {"name": self.name, "sensors": self.sensors}
